import os
import json
import time
import gradio as gr
from datetime import datetime
from openai import OpenAI
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize OpenAI client
client = OpenAI(
    api_key=os.getenv("LLM_API_KEY"),
    base_url=os.getenv("LLM_BASE_URL"),
)

# Global variable for history
story_history = []

def load_history():
    """加载历史记录"""
    global story_history
    try:
        if os.path.exists("story_history.json"):
            with open("story_history.json", "r", encoding="utf-8") as f:
                story_history = json.load(f)
    except:
        story_history = []

def save_history():
    """保存历史记录"""
    global story_history
    try:
        with open("story_history.json", "w", encoding="utf-8") as f:
            json.dump(story_history, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"保存历史记录失败: {e}")

def generate_story(theme, keywords, genre, length, language):
    genre_descriptions = {
        "奇幻": "fantasy with magic and mythical creatures",
        "科幻": "science fiction with advanced technology and futuristic elements",
        "现实": "realistic contemporary story",
        "悬疑": "mystery or thriller with suspense",
        "爱情": "romantic story focused on relationships",
        "冒险": "adventure with exploration and excitement"
    }

    genre_desc = genre_descriptions.get(genre, "general fiction")
    prompt = f"""Write a {length} story in {language} about {theme}.
Genre: {genre_desc}.
Include keywords: {keywords}.
Make the story engaging and well-structured.
"""

    try:
        max_tokens = 1500 if "短篇" in length else 3000
        response = client.chat.completions.create(
            model=os.getenv("LLM_MODEL"),
            messages=[{"role": "user", "content": prompt}],
            max_tokens=max_tokens,
        )
        story = response.choices[0].message.content.strip()

        # 保存到历史记录
        story_entry = {
            "theme": theme,
            "keywords": keywords,
            "genre": genre,
            "length": length,
            "language": language,
            "story": story[:200] + "..." if len(story) > 200 else story,
            "timestamp": time.time()
        }
        story_history.insert(0, story_entry)
        if len(story_history) > 20:  # 限制历史记录数量
            story_history.pop()
        save_history()

        return story
    except Exception as e:
        return f"生成失败，请检查网络连接或稍后重试。\n错误详情: {str(e)}"

def get_history_display():
    """获取历史记录的显示格式"""
    load_history()
    if not story_history:
        return "暂无历史记录"

    display_text = []
    for i, entry in enumerate(story_history[:10], 1):  # 显示前10条
        timestamp = datetime.fromtimestamp(entry.get('timestamp', 0)).strftime('%Y-%m-%d %H:%M:%S')
        display_text.append(f"{i}. [{timestamp}] 主题: {entry.get('theme', '')} | 类型: {entry.get('genre', '')} | 语言: {entry.get('language', '')}\n   {entry.get('story', '')}\n")

    return "\n".join(display_text)

# Custom CSS for better UI
css = """
.gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    font-family: 'Arial', sans-serif;
}
.gr-button {
    background: linear-gradient(45deg, #ff6b35, #f7931e);
    border: none;
    border-radius: 25px;
    color: white;
    font-weight: bold;
}
.gr-button:hover {
    background: linear-gradient(45deg, #f7931e, #ff6b35);
    transform: scale(1.05);
}
.gr-input {
    border-radius: 10px;
    border: 2px solid #ddd;
}
.gr-textbox {
    border-radius: 10px;
    border: 2px solid #ddd;
}
"""

# Gradio interface with improved UI and tabs
with gr.Blocks(title="AI 故事生成器", css=css, theme=gr.themes.Soft()) as iface:
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 20px;">
        <h1 style="color: white; text-shadow: 2px 2px 4px rgba(0,0,0,0.3);">
            🪄 AI 创意故事生成器
        </h1>
        <p style="color: rgba(255,255,255,0.8); font-size: 16px;">
            让 AI 为您创作精彩的故事
        </p>
    </div>
    """)

    with gr.Tabs():
        with gr.TabItem("📝 生成故事"):
            with gr.Row():
                with gr.Column(scale=2):
                    theme_input = gr.Textbox(
                        label="📚 故事主题",
                        placeholder="输入故事主题，如冒险、爱情、友情等",
                        lines=2
                    )

                    keywords_input = gr.Textbox(
                        label="🔑 关键词",
                        placeholder="输入关键词，用逗号分隔，如勇气、魔法、城市",
                        lines=2
                    )

                    with gr.Row():
                        genre_input = gr.Dropdown(
                            label="🎭 故事类型",
                            choices=["奇幻", "科幻", "现实", "悬疑", "爱情", "冒险"],
                            value="现实"
                        )

                        length_input = gr.Dropdown(
                            label="📏 故事长度",
                            choices=["短篇 (1500字以内)", "中篇 (3000字以内)"],
                            value="短篇 (1500字以内)"
                        )

                    language_input = gr.Radio(
                        label="🌍 输出语言",
                        choices=["中文", "English"],
                        value="中文"
                    )

                    with gr.Row():
                        generate_btn = gr.Button(
                            "🚀 生成故事",
                            variant="primary",
                            size="lg"
                        )

                with gr.Column(scale=3):
                    output_textbox = gr.Textbox(
                        label="📖 生成的故事",
                        lines=15,
                        show_copy_button=True
                    )

                    status_box = gr.Textbox(
                        label="📢 状态信息",
                        interactive=False,
                        visible=False
                    )

        with gr.TabItem("📚 历史记录"):
            history_textbox = gr.Textbox(
                label="📖 历史故事记录",
                lines=20,
                interactive=False,
                value=get_history_display()
            )

            with gr.Row():
                refresh_history_btn = gr.Button("🔄 刷新记录", size="sm")
                clear_history_btn = gr.Button("🗑️ 清除历史", size="sm")

    # Event handlers
    def generate_with_status(theme, keywords, genre, length, language, status):
        # 显示加载状态
        status_update = gr.update(value="正在生成故事，请稍候...", visible=True)
        story = generate_story(theme, keywords, genre, length, language)
        # 完成后隐藏状态
        final_status = gr.update(visible=False)
        return story, final_status

    generate_btn.click(
        fn=generate_with_status,
        inputs=[theme_input, keywords_input, genre_input, length_input, language_input, status_box],
        outputs=[output_textbox, status_box]
    )

    refresh_history_btn.click(
        fn=lambda: get_history_display(),
        inputs=[],
        outputs=history_textbox
    )

    def clear_history():
        global story_history
        story_history = []
        save_history()
        return "历史记录已清除", get_history_display()

    clear_history_btn.click(
        fn=clear_history,
        inputs=[],
        outputs=[status_box, history_textbox]
    )

if __name__ == "__main__":
    iface.launch(
        share=True,
        show_error=True,
        server_name="0.0.0.0",
        server_port=7860
    )
